"""Contains a variant of the densenet model definition."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

slim = tf.contrib.slim


def trunc_normal(stddev): return tf.truncated_normal_initializer(stddev=stddev)

#这个代码是块中每个小的操作,dropout 扩展了模型的通用性，泛化性
def bn_act_conv_drp(current, num_outputs, kernel_size, scope='block'):
    current = slim.batch_norm(current, scope=scope + '_bn')
    current = tf.nn.relu(current)
    current = slim.conv2d(current, num_outputs, kernel_size, scope=scope + '_conv')
    current = slim.dropout(current, scope=scope + '_dropout')
    return current

#这个代码是DenseBlock
def block(net, layers, growth, scope='block'):
    for idx in range(layers):
        bottleneck = bn_act_conv_drp(net, 4 * growth, [1, 1],
                                     scope=scope + '_conv1x1' + str(idx))
        tmp = bn_act_conv_drp(bottleneck, growth, [3, 3],
                              scope=scope + '_conv3x3' + str(idx))
        net = tf.concat(axis=3, values=[net, tmp])
    return net


def densenet(images, num_classes=1001, is_training=False,
             dropout_keep_prob=0.8,
             scope='densenet'):
    """Creates a variant of the densenet model.

      images: A batch of `Tensors` of size [batch_size, height, width, channels].
      num_classes: the number of classes in the dataset.
      is_training: specifies whether or not we're currently training the model.
        This variable will determine the behaviour of the dropout layer.
      dropout_keep_prob: the percentage of activation values that are retained.
      prediction_fn: a function to get predictions out of logits.
      scope: Optional variable_scope.

    Returns:
      logits: the pre-softmax activations, a tensor of size
        [batch_size, `num_classes`]
      end_points: a dictionary from components of the network to the corresponding
        activation.
    """
    #config = tf.ConfigProto(allow_soft_placement=True)
    growth = 12
    compression_rate = 0.5

    def reduce_dim(input_feature):
        return int(int(input_feature.shape[-1]) * compression_rate)

    end_points = {}

    with tf.variable_scope(scope, 'DenseNet', [images, num_classes]):
        with slim.arg_scope(bn_drp_scope(is_training=is_training,
                                         keep_prob=dropout_keep_prob)) as ssc:
            pass
            ##########################
            # Put your code here.
            ##########################
            print("@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@")
            
            #1.init
            net = slim.conv2d(images, 2 * growth, [7, 7], stride=2,activation_fn=None, padding='same', scope='init')
            net = slim.max_pool2d(net, [3, 3], stride=2, scope="init_pool")
            
            #2.DenseBlock, 进行12次小block， 每个小block包括bn ，relu， 卷积3*3，输出为 growth， 如果growth太大导致输入增长太快，运行慢
			#growth太小， 可能有些信息没有提取出来
            net = block(net, 12, growth, scope='block1')
            
            #3.两block之间的卷积和池化，变换层，1.瓶颈层： 1*1 降维，BN_Relu_Conv。。。。 DenseNet-B
            net = bn_act_conv_drp(net, growth, [1, 1], scope='con2d1')
            #net = slim.conv2d(net, growth, [3,3], scope="con2d1")
            net = slim.avg_pool2d(net, [2,2], stride=2, scope="pool1")
            
            #2.DenseBlock
            net = block(net, 12, growth, scope='block2')
            
            #3.两block之间的卷积和池化
            net = bn_act_conv_drp(net, growth, [1, 1], scope='con2d2')
            net = slim.avg_pool2d(net, [2,2], stride=2, scope="pool2")
            
            #2.DenseBlock
            net = block(net, 12, growth, scope='block3')
            
            #3.两block之间的卷积和池化
            net = bn_act_conv_drp(net, growth, [1, 1], scope='con2d3')
            net = slim.avg_pool2d(net, [2,2], stride=2, scope="pool3")
            
            
            #2.DenseBlock
            net = block(net, 12, growth, scope='block4')
            
            #3.全局池化
            net = slim.avg_pool2d(net, [6,6], stride=1, scope="pool4")
            
            net = slim.flatten(net, scope='Flatten')
            
            
            logits = slim.fully_connected(net, num_classes, activation_fn=None, scope="logits")
            end_points['softmax'] = tf.nn.softmax(logits, name="Predictions")
            
    print("@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@")

    return logits, end_points


def bn_drp_scope(is_training=True, keep_prob=0.8):
    keep_prob = keep_prob if is_training else 1
    with slim.arg_scope(
        [slim.batch_norm],
            scale=True, is_training=is_training, updates_collections=None):
        with slim.arg_scope(
            [slim.dropout],
                is_training=is_training, keep_prob=keep_prob) as bsc:
            return bsc


def densenet_arg_scope(weight_decay=0.004):
    """Defines the default densenet argument scope.

    Args:
      weight_decay: The weight decay to use for regularizing the model.

    Returns:
      An `arg_scope` to use for the inception v3 model.
    """
    with slim.arg_scope(
        [slim.conv2d],
        weights_initializer=tf.contrib.layers.variance_scaling_initializer(
            factor=2.0, mode='FAN_IN', uniform=False),
        activation_fn=None, biases_initializer=None, padding='same',
            stride=1) as sc:
        return sc


densenet.default_image_size = 224
"""
最后做下总结：
1.growth， 如果growth太小，那么会导致 一些信息提取不错来
太大 又会导致 运行满，内存消耗等问题，
这次作业没有运行特别长时间， 但是通过比较 growth = 12， 24， 32 发现：
在12时 其实运行最快，并且最后输出的loss 也下降的较快， 看来是比较好的
而 24 和 32  比较， loss都比较大，并且， 32运行太慢
2.关于densenet的一些理解
a.拼接之前层的输出， 可以看作是 第一层对输入数据经行简单提取特征， 然后交给第二层，
 而第二层 不仅在第一层的输出上进行学习， 还结合第一层的输入来学习，
 有点类似adaboost，在前一层的学习结果上，不断使用原来的数据，再次学习
 或者 也有可以认为，第一层学习后， 如果交给第二层直接学习，那么会有一些信息损失，
 如果继续结合第一层的输入，就会把损失的数据补充回来，达到更好的学习效果
b.由于不断的拼接前面的输出数据，稠密连接，出现的后果， 通道数增长太快，论文中使用1*1卷积降维（瓶颈层），
 dropout 实际 也 提高了 网络的效率，
 但有些可能是没有用的，比如对某一层 可能 前面的第一层没有用，但是无法保证 第一层对最后一层没有用，所以只能交给模型来自己选择
c.bn，relu， con2v 为什么会作为一个整体呢？bn类似数据标准化，其实不会对结果产生太大影响，relu也只是把结果稍作修改，基本保持性质不变
真正起作用的是 con2v， 所以把他们作为一个整体来处理 能使运行更简化
"""
